import json
import os
import random
import urllib
from tqdm import tqdm

# CHANGES:

# + llm_api.py

# lm_eval/base.py line:516
# + trust_remote_code=True,

# main.py 
# + import llm_api
# + lm_eval.models.MODEL_REGISTRY['llm_api'] = llm_api.LLM_API
# -/+ parser.add_argument("--model", default="llm_api")
# -/+ parser.add_argument("--model_args", default="engine=llama3_8B")
# -/+ parser.add_argument("--tasks", default='gsm8k', choices=utils.MultiChoice(tasks.ALL_TASKS))
# -/+ parser.add_argument("--num_fewshot", type=int, default=3)

'''
export HF_ENDPOINT="https://hf-mirror.com"
export HF_DATASETS_CACHE="datasets_cache"
export NUMEXPR_MAX_THREADS=64
'''

# To solve: <urlopen error [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: unable to get local issuer certificate (_ssl.c:1000)>
# ln -s /etc/ssl/* /Library/Frameworks/Python.framework/Versions/3.12/etc/openssl

# def list_tasks():
#     from lm_eval import tasks
#     task_list = [(i, task_name) for i, (task_name, task_cls) in enumerate(tasks.TASK_REGISTRY.items())]
#     print('task_list = [')
#     for i,tn in task_list:
#         print(f'\t"{tn}", #{i+1}')
#     print('\t]')

task_list = [
    "cola", #1
    "mnli", #2
    "mnli_mismatched", #3
    "mrpc", #4
    "rte", #5
    "qnli", #6
    "qqp", #7
    "sst", #8
    "wnli", #9
    "boolq", #10
    "cb", #11
    "copa", #12
    "multirc", #13
    "record", #14
    "wic", #15
    "wsc", #16
    "coqa", #17
    "drop", #18
    "lambada_openai", #19
    "lambada_standard", #20
    "lambada_openai_cloze", #21
    "lambada_standard_cloze", #22
    "lambada_openai_mt_en", #23
    "lambada_openai_mt_fr", #24
    "lambada_openai_mt_de", #25
    "lambada_openai_mt_it", #26
    "lambada_openai_mt_es", #27
    "wikitext", #28
    "piqa", #29
    "prost", #30
    "mc_taco", #31
    "pubmedqa", #32
    "sciq", #33
    "qasper", #34
    "qa4mre_2011", #35
    "qa4mre_2012", #36
    "qa4mre_2013", #37
    "triviaqa", #38
    "arc_easy", #39
    "arc_challenge", #40
    "logiqa", #41
    "hellaswag", #42
    "swag", #43
    "openbookqa", #44
    "squad2", #45
    "race", #46
    "headqa", #47
    "headqa_es", #48
    "headqa_en", #49
    "mathqa", #50
    "webqs", #51
    "wsc273", #52
    "winogrande", #53
    "anli_r1", #54
    "anli_r2", #55
    "anli_r3", #56
    "ethics_cm", #57
    "ethics_deontology", #58
    "ethics_justice", #59
    "ethics_utilitarianism_original", #60
    "ethics_utilitarianism", #61
    "ethics_virtue", #62
    "truthfulqa_mc", #63
    "truthfulqa_gen", #64
    "mutual", #65
    "mutual_plus", #66
    "math_algebra", #67
    "math_counting_and_prob", #68
    "math_geometry", #69
    "math_intermediate_algebra", #70
    "math_num_theory", #71
    "math_prealgebra", #72
    "math_precalc", #73
    "math_asdiv", #74
    "gsm8k", #75
    "arithmetic_2da", #76
    "arithmetic_2ds", #77
    "arithmetic_3da", #78
    "arithmetic_3ds", #79
    "arithmetic_4da", #80
    "arithmetic_4ds", #81
    "arithmetic_5da", #82
    "arithmetic_5ds", #83
    "arithmetic_2dm", #84
    "arithmetic_1dc", #85
    "hendrycksTest-abstract_algebra", #86
    "hendrycksTest-anatomy", #87
    "hendrycksTest-astronomy", #88
    "hendrycksTest-business_ethics", #89
    "hendrycksTest-clinical_knowledge", #90
    "hendrycksTest-college_biology", #91
    "hendrycksTest-college_chemistry", #92
    "hendrycksTest-college_computer_science", #93
    "hendrycksTest-college_mathematics", #94
    "hendrycksTest-college_medicine", #95
    "hendrycksTest-college_physics", #96
    "hendrycksTest-computer_security", #97
    "hendrycksTest-conceptual_physics", #98
    "hendrycksTest-econometrics", #99
    "hendrycksTest-electrical_engineering", #100
    "hendrycksTest-elementary_mathematics", #101
    "hendrycksTest-formal_logic", #102
    "hendrycksTest-global_facts", #103
    "hendrycksTest-high_school_biology", #104
    "hendrycksTest-high_school_chemistry", #105
    "hendrycksTest-high_school_computer_science", #106
    "hendrycksTest-high_school_european_history", #107
    "hendrycksTest-high_school_geography", #108
    "hendrycksTest-high_school_government_and_politics", #109
    "hendrycksTest-high_school_macroeconomics", #110
    "hendrycksTest-high_school_mathematics", #111
    "hendrycksTest-high_school_microeconomics", #112
    "hendrycksTest-high_school_physics", #113
    "hendrycksTest-high_school_psychology", #114
    "hendrycksTest-high_school_statistics", #115
    "hendrycksTest-high_school_us_history", #116
    "hendrycksTest-high_school_world_history", #117
    "hendrycksTest-human_aging", #118
    "hendrycksTest-human_sexuality", #119
    "hendrycksTest-international_law", #120
    "hendrycksTest-jurisprudence", #121
    "hendrycksTest-logical_fallacies", #122
    "hendrycksTest-machine_learning", #123
    "hendrycksTest-management", #124
    "hendrycksTest-marketing", #125
    "hendrycksTest-medical_genetics", #126
    "hendrycksTest-miscellaneous", #127
    "hendrycksTest-moral_disputes", #128
    "hendrycksTest-moral_scenarios", #129
    "hendrycksTest-nutrition", #130
    "hendrycksTest-philosophy", #131
    "hendrycksTest-prehistory", #132
    "hendrycksTest-professional_accounting", #133
    "hendrycksTest-professional_law", #134
    "hendrycksTest-professional_medicine", #135
    "hendrycksTest-professional_psychology", #136
    "hendrycksTest-public_relations", #137
    "hendrycksTest-security_studies", #138
    "hendrycksTest-sociology", #139
    "hendrycksTest-us_foreign_policy", #140
    "hendrycksTest-virology", #141
    "hendrycksTest-world_religions", #142
    "wmt14-en-fr", #143
    "wmt14-fr-en", #144
    "wmt16-en-ro", #145
    "wmt16-ro-en", #146
    "wmt16-de-en", #147
    "wmt16-en-de", #148
    "wmt20-cs-en", #149
    "wmt20-de-en", #150
    "wmt20-de-fr", #151
    "wmt20-en-cs", #152
    "wmt20-en-de", #153
    "wmt20-en-iu", #154
    "wmt20-en-ja", #155
    "wmt20-en-km", #156
    "wmt20-en-pl", #157
    "wmt20-en-ps", #158
    "wmt20-en-ru", #159
    "wmt20-en-ta", #160
    "wmt20-en-zh", #161
    "wmt20-fr-de", #162
    "wmt20-iu-en", #163
    "wmt20-ja-en", #164
    "wmt20-km-en", #165
    "wmt20-pl-en", #166
    "wmt20-ps-en", #167
    "wmt20-ru-en", #168
    "wmt20-ta-en", #169
    "wmt20-zh-en", #170
    "iwslt17-en-ar", #171
    "iwslt17-ar-en", #172
    "anagrams1", #173
    "anagrams2", #174
    "cycle_letters", #175
    "random_insertion", #176
    "reversed_words", #177
    "pile_arxiv", #178
    "pile_books3", #179
    "pile_bookcorpus2", #180
    "pile_dm-mathematics", #181
    "pile_enron", #182
    "pile_europarl", #183
    "pile_freelaw", #184
    "pile_github", #185
    "pile_gutenberg", #186
    "pile_hackernews", #187
    "pile_nih-exporter", #188
    "pile_opensubtitles", #189
    "pile_openwebtext2", #190
    "pile_philpapers", #191
    "pile_pile-cc", #192
    "pile_pubmed-abstracts", #193
    "pile_pubmed-central", #194
    "pile_stackexchange", #195
    "pile_uspto", #196
    "pile_ubuntu-irc", #197
    "pile_wikipedia", #198
    "pile_youtubesubtitles", #199
    "blimp_adjunct_island", #200
    "blimp_anaphor_gender_agreement", #201
    "blimp_anaphor_number_agreement", #202
    "blimp_animate_subject_passive", #203
    "blimp_animate_subject_trans", #204
    "blimp_causative", #205
    "blimp_complex_NP_island", #206
    "blimp_coordinate_structure_constraint_complex_left_branch", #207
    "blimp_coordinate_structure_constraint_object_extraction", #208
    "blimp_determiner_noun_agreement_1", #209
    "blimp_determiner_noun_agreement_2", #210
    "blimp_determiner_noun_agreement_irregular_1", #211
    "blimp_determiner_noun_agreement_irregular_2", #212
    "blimp_determiner_noun_agreement_with_adj_2", #213
    "blimp_determiner_noun_agreement_with_adj_irregular_1", #214
    "blimp_determiner_noun_agreement_with_adj_irregular_2", #215
    "blimp_determiner_noun_agreement_with_adjective_1", #216
    "blimp_distractor_agreement_relational_noun", #217
    "blimp_distractor_agreement_relative_clause", #218
    "blimp_drop_argument", #219
    "blimp_ellipsis_n_bar_1", #220
    "blimp_ellipsis_n_bar_2", #221
    "blimp_existential_there_object_raising", #222
    "blimp_existential_there_quantifiers_1", #223
    "blimp_existential_there_quantifiers_2", #224
    "blimp_existential_there_subject_raising", #225
    "blimp_expletive_it_object_raising", #226
    "blimp_inchoative", #227
    "blimp_intransitive", #228
    "blimp_irregular_past_participle_adjectives", #229
    "blimp_irregular_past_participle_verbs", #230
    "blimp_irregular_plural_subject_verb_agreement_1", #231
    "blimp_irregular_plural_subject_verb_agreement_2", #232
    "blimp_left_branch_island_echo_question", #233
    "blimp_left_branch_island_simple_question", #234
    "blimp_matrix_question_npi_licensor_present", #235
    "blimp_npi_present_1", #236
    "blimp_npi_present_2", #237
    "blimp_only_npi_licensor_present", #238
    "blimp_only_npi_scope", #239
    "blimp_passive_1", #240
    "blimp_passive_2", #241
    "blimp_principle_A_c_command", #242
    "blimp_principle_A_case_1", #243
    "blimp_principle_A_case_2", #244
    "blimp_principle_A_domain_1", #245
    "blimp_principle_A_domain_2", #246
    "blimp_principle_A_domain_3", #247
    "blimp_principle_A_reconstruction", #248
    "blimp_regular_plural_subject_verb_agreement_1", #249
    "blimp_regular_plural_subject_verb_agreement_2", #250
    "blimp_sentential_negation_npi_licensor_present", #251
    "blimp_sentential_negation_npi_scope", #252
    "blimp_sentential_subject_island", #253
    "blimp_superlative_quantifiers_1", #254
    "blimp_superlative_quantifiers_2", #255
    "blimp_tough_vs_raising_1", #256
    "blimp_tough_vs_raising_2", #257
    "blimp_transitive", #258
    "blimp_wh_island", #259
    "blimp_wh_questions_object_gap", #260
    "blimp_wh_questions_subject_gap", #261
    "blimp_wh_questions_subject_gap_long_distance", #262
    "blimp_wh_vs_that_no_gap", #263
    "blimp_wh_vs_that_no_gap_long_distance", #264
    "blimp_wh_vs_that_with_gap", #265
    "blimp_wh_vs_that_with_gap_long_distance", #266
    "toxigen", #267
    "crows_pairs_english", #268
    "crows_pairs_english_race_color", #269
    "crows_pairs_english_socioeconomic", #270
    "crows_pairs_english_gender", #271
    "crows_pairs_english_age", #272
    "crows_pairs_english_religion", #273
    "crows_pairs_english_disability", #274
    "crows_pairs_english_sexual_orientation", #275
    "crows_pairs_english_nationality", #276
    "crows_pairs_english_physical_appearance", #277
    "crows_pairs_english_autre", #278
    "crows_pairs_french", #279
    "crows_pairs_french_race_color", #280
    "crows_pairs_french_socioeconomic", #281
    "crows_pairs_french_gender", #282
    "crows_pairs_french_age", #283
    "crows_pairs_french_religion", #284
    "crows_pairs_french_disability", #285
    "crows_pairs_french_sexual_orientation", #286
    "crows_pairs_french_nationality", #287
    "crows_pairs_french_physical_appearance", #288
    "crows_pairs_french_autre", #289
    "xcopa_et", #290
    "xcopa_ht", #291
    "xcopa_it", #292
    "xcopa_id", #293
    "xcopa_qu", #294
    "xcopa_sw", #295
    "xcopa_zh", #296
    "xcopa_ta", #297
    "xcopa_th", #298
    "xcopa_tr", #299
    "xcopa_vi", #300
    "bigbench_temporal_sequences", #301
    "bigbench_tracking_shuffled_objects_five_objects", #302
    "bigbench_salient_translation_error_detection", #303
    "bigbench_snarks", #304
    "bigbench_ruin_names", #305
    "bigbench_navigate", #306
    "bigbench_date_understanding", #307
    "bigbench_hyperbaton", #308
    "bigbench_dyck_languages", #309
    "bigbench_tracking_shuffled_objects_three_objects", #310
    "bigbench_tracking_shuffled_objects_seven_objects", #311
    "bigbench_causal_judgement", #312
    "bigbench_sports_understanding", #313
    "bigbench_logical_deduction_five_objects", #314
    "bigbench_movie_recommendation", #315
    "bigbench_logical_deduction_three_objects", #316
    "bigbench_geometric_shapes", #317
    "bigbench_formal_fallacies_syllogisms_negation", #318
    "bigbench_disambiguation_qa", #319
    "bigbench_logical_deduction_seven_objects", #320
    "bigbench_reasoning_about_colored_objects", #321
    "xstory_cloze_en", #322
    "xstory_cloze_ru", #323
    "xstory_cloze_zh", #324
    "xstory_cloze_es", #325
    "xstory_cloze_ar", #326
    "xstory_cloze_hi", #327
    "xstory_cloze_id", #328
    "xstory_cloze_te", #329
    "xstory_cloze_sw", #330
    "xstory_cloze_eu", #331
    "xstory_cloze_my", #332
    "xwinograd_en", #333
    "xwinograd_fr", #334
    "xwinograd_jp", #335
    "xwinograd_pt", #336
    "xwinograd_ru", #337
    "xwinograd_zh", #338
    "pawsx_en", #339
    "pawsx_de", #340
    "pawsx_es", #341
    "pawsx_fr", #342
    "pawsx_ja", #343
    "pawsx_ko", #344
    "pawsx_zh", #345
    "xnli_ar", #346
    "xnli_bg", #347
    "xnli_de", #348
    "xnli_el", #349
    "xnli_en", #350
    "xnli_es", #351
    "xnli_fr", #352
    "xnli_hi", #353
    "xnli_ru", #354
    "xnli_sw", #355
    "xnli_th", #356
    "xnli_tr", #357
    "xnli_ur", #358
    "xnli_vi", #359
    "xnli_zh", #360
    "mgsm_en", #361
    "mgsm_es", #362
    "mgsm_fr", #363
    "mgsm_de", #364
    "mgsm_ru", #365
    "mgsm_zh", #366
    "mgsm_ja", #367
    "mgsm_th", #368
    "mgsm_sw", #369
    "mgsm_bn", #370
    "mgsm_te", #371
    ]
    

def info(start=0,end=1000):
    from lm_eval import tasks, evaluator, utils
    from lm_eval.base import PerplexityTask
    import datasets
    exceptions = (
        ConnectionError,
        urllib.error.URLError,
        FileNotFoundError,
        datasets.exceptions.DatasetGenerationError,
        )

    rnd = random.Random()
    total = len(tasks.TASK_REGISTRY)
    count = 0
    info = []
    for task_name, task_cls in tqdm(tasks.TASK_REGISTRY.items(), total=total):
        count += 1
        if count not in range(start,end): 
            continue
        print(f'# {count}/{total}')
        cls_name = f'{task_cls.__module__}.{task_cls.__name__}'
        print(task_name, cls_name)
        try:
            task = task_cls()
        except exceptions as ex:
            import traceback
            traceback.print_exc()
            info.append({'task':task_name, 'class':cls_name,
                        'exception': str(ex)})
            continue
        task_doc_func = task.test_docs if task.has_test_docs() else task.validation_docs
        task_docs = list(task_doc_func())
        doc = task_docs[0]
        num_fewshot = 3
        for _ in range(2):
            try:
                ctx = task.fewshot_context(doc=doc, num_fewshot=num_fewshot, 
                                        rnd=rnd, description='')
                break
            except AssertionError:
                num_fewshot = 0
        print('\tnum_task_docs:', len(task_docs))
        print('\tnum_fewshot:', num_fewshot)
        reqs = task.construct_requests(doc, ctx)
        if not isinstance(reqs, (list,tuple)):
            reqs = [reqs]
        request_type = []
        for req in reqs:
            print('\trequest_type:', req.request_type)
            request_type.append(req.request_type)
        info.append({'task':task_name, 'class':cls_name,
                    'size':len(task_docs), 'num_fewshot':num_fewshot,
                    'request_type':request_type})
    with open('exam_tasks_info.json', 'w') as fp:
        json.dump(info, fp, indent='\t')

def exam():
    with open('exam_tasks_info.json') as fp:
        tasks = json.load(fp)
    exceptions = {}
    request_type_count = {}
    greedy_until_tasks = []
    for task in tasks:
        if 'exception' in task:
            if task['exception'] not in exceptions:
                exceptions[task['exception']] = []
            exceptions[task['exception']].append(task['task'])
        else:
            if len(set(task['request_type']))!=1:
                continue
            request_type = task['request_type'][0]
            request_type_count[request_type] = 1 + request_type_count.get(request_type,0)
            if request_type=='greedy_until':
                greedy_until_tasks.append(task['task'])

    print('exceptions = ', end='')
    print(json.dumps(exceptions, indent='\t'))
    print('request_type_count = ', end='')
    print(json.dumps(request_type_count, indent='\t'))
    print('greedy_until_tasks = ', end='')
    print(json.dumps(greedy_until_tasks, indent='\t'))


# An error occurred while generating the dataset
# Tasks: headqa, headqa_es, headqa_en

# Couldn't find file at https://the-eye.eu/public/AI/pile/val.jsonl.zst
# Tasks: pile_arxiv, pile_books3, pile_bookcorpus2, pile_dm-mathematics, pile_enron, pile_europarl, pile_freelaw, pile_github, pile_gutenberg, pile_hackernews, pile_nih-exporter, pile_opensubtitles, pile_openwebtext2, pile_philpapers, pile_pile-cc, pile_pubmed-abstracts, pile_pubmed-central, pile_stackexchange, pile_uspto, pile_ubuntu-irc, pile_wikipedia, pile_youtubesubtitles

# request_type_count:
# - greedy_until : 28
# - loglikelihood : 262
# - loglikelihood_rolling : 23
greedy_until_tasks = [
    "coqa",
    "drop",
    "triviaqa",
    "truthfulqa_gen",
    "math_algebra",
    "math_counting_and_prob",
    "math_geometry",
    "math_intermediate_algebra",
    "math_num_theory",
    "math_prealgebra",
    "math_precalc",
    "gsm8k",
    "anagrams1",
    "anagrams2",
    "cycle_letters",
    "random_insertion",
    "reversed_words",
    "mgsm_en",
    "mgsm_es",
    "mgsm_fr",
    "mgsm_de",
    "mgsm_ru",
    "mgsm_zh",
    "mgsm_ja",
    "mgsm_th",
    "mgsm_sw",
    "mgsm_bn",
    "mgsm_te"
]

if __name__ == "__main__":
    info()
    exam()
